import pandas as pd
from pymilvus import (
    connections, FieldSchema, CollectionSchema, DataType,
    Collection
)
from langchain.schema import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import JinaEmbeddings


JINA_API_KEY = "jina_0806183d486d4057a8822c3f11cc59adq36ajgW3bQbNE7AT5m9qANB1uOye"
EMBEDDING_DIM = 1024


embedding_model = JinaEmbeddings(
    jina_api_key=JINA_API_KEY,
    model_name="jina-embeddings-v3"
)


herb_df = pd.read_csv("MedicineTable.csv", encoding="utf-8")

#构造Document对象
herb_documents = []
for _, row in herb_df.iterrows():
    content_parts = [
        str(row["名字"]) if pd.notna(row["名字"]) else "",
        str(row["别名"]) if pd.notna(row["别名"]) else "",
        str(row["主治症状"]) if pd.notna(row["主治症状"]) else "",
        str(row["性味"]) if pd.notna(row["性味"]) else "",
        str(row["归经"]) if pd.notna(row["归经"]) else "",
        str(row["产地"]) if pd.notna(row["产地"]) else "",
        str(row["来源"]) if pd.notna(row["来源"]) else ""
    ]
    full_text = "；".join(filter(None, content_parts))
    metadata = {"id": int(row["id"]), "name": row["名字"]}
    herb_documents.append(Document(page_content=full_text, metadata=metadata))

#切分文档
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
split_documents = text_splitter.split_documents(herb_documents)

#提取文本
texts = [doc.page_content for doc in split_documents]
metadatas = [doc.metadata for doc in split_documents]


print("正在生成向量...")
embeddings = embedding_model.embed_documents(texts)


connections.connect(alias="default", host="localhost", port="19530")


collection_name = "medicine"
id_field = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True)
vector_field = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=EMBEDDING_DIM)
text_field = FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535)
metadata_field = FieldSchema(name="metadata", dtype=DataType.JSON)

schema = CollectionSchema(
    fields=[id_field, vector_field, text_field, metadata_field],
    description="Medicine schema"
)

collection = Collection(name=collection_name, schema=schema)


print("正在写入 Milvus 向量数据库...")
collection.insert([embeddings, texts, metadatas])

#创建索引
index_params = {
    "index_type": "IVF_FLAT",
    "metric_type": "L2",
    "params": {"nlist": 128}
}
collection.create_index(field_name="vector", index_params=index_params)

print("向量生成与写入完成")
